import pandas as pd
import os

# 定义标签映射
label_mapping = {
    'E3TF': 'E3TF_1',
    'E3TF普通': 'PT',
    'P8普通': 'PT',
    'P8_椭圆': 'P8_1',
    'P8_跑道': 'P8_2',
    'R135': 'R135_1',
    'R135普通': 'PT'
}

# 定义训练数据目录
train_data_directory = r'D:\pythonProject\模式识别1115\训练集_备份'  # 替换为您的目录路径

# 遍历训练数据目录下的所有子文件夹
for foldername in os.listdir(train_data_directory):
    folder_path = os.path.join(train_data_directory, foldername)

    # 确保是目录
    if os.path.isdir(folder_path):
        # 获取对应的标签
        label = label_mapping.get(foldername)

        if label is not None:
            # 初始化文件计数器
            file_count = 1

            # 遍历当前子文件夹下的所有 CSV 文件
            for filename in os.listdir(folder_path):
                if filename.endswith('.csv'):
                    # 读取 CSV 文件
                    file_path = os.path.join(folder_path, filename)
                    df = pd.read_csv(file_path)

                    # 添加标签列
                    df['Label'] = label

                    # 生成新的文件名
                    new_filename = f"{label}({file_count}).csv"
                    new_file_path = os.path.join(train_data_directory, new_filename)  # 直接保存在train_data_directory

                    # 保存回新的 CSV 文件
                    df.to_csv(new_file_path, index=False)
                    print(f'Saved {new_filename} in {train_data_directory} with label {label}.')

                    # 更新文件计数器
                    file_count += 1
        else:
            print(f'No matching label for folder: {foldername}.')